import os
import torch
import random
import numpy as np

from data import Data, DataPlus
from experiment import Experiment
from model import Learner
from opts import parse_opt


def set_seed(opt, n_gpu):
    random.seed(opt.seed)
    np.random.seed(opt.seed)
    torch.manual_seed(opt.seed)
    if n_gpu > 0:
        torch.cuda.manual_seed_all(opt.seed)

def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    n_gpu = torch.cuda.device_count()

    option = parse_opt()

    set_seed(option, n_gpu)

    if not option.query_is_language:
        data = Data(option.datadir, option.seed, option.type_check,
                    option.domain_size, option.no_extra_facts)
    else:
        data = DataPlus(option.datadir, option.seed)
    print("Data prepared.")

    option.num_entity = data.num_entity
    option.num_operator = data.num_operator
    if not option.query_is_language:
        option.num_query = data.num_query
    else:
        option.num_vocab = data.num_vocab
        option.num_word = data.num_word  # the number of words in each query

    option.this_expsdir = os.path.join(option.exps_dir, option.tag)
    if not os.path.exists(option.this_expsdir):
        os.makedirs(option.this_expsdir)
    option.ckpt_dir = os.path.join(option.this_expsdir, "ckpt")
    if not os.path.exists(option.ckpt_dir):
        os.makedirs(option.ckpt_dir)
    option.model_path = os.path.join(option.ckpt_dir, "model")

    option.save()
    print("Option saved.")

    learner = Learner(option, device)
    learner.to(device)
    print("Learner built.")

    if option.from_model_ckpt is not None:
        pass

    data.reset(option.batch_size)
    experiment = Experiment(option, learner, data)
    print("Experiment created.")

    if not option.no_train:
        print("Start training...")
        experiment.train()

    if not option.no_preds:
        print("Start getting test predictions...")
        experiment.get_predictions()

    if not option.no_rules:
        print("Start getting rules...")
        experiment.get_rules()

    if option.get_vocab_embed:
        print("Start getting vocabulary embedding...")
        experiment.get_vocab_embedding()

    experiment.close_log_file()
    print("=" * 36 + "Finish" + "=" * 36)

if __name__ == '__main__':
    main()